/*
 * AIEngine a new generation network intrusion detection system.
 *
 * Copyright (C) 2013-2023  Luis Campo Giralte
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Library General Public
 * License as published by the Free Software Foundation; either
 * version 2 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Library General Public License for more details.
 *
 * You should have received a copy of the GNU Library General Public
 * License along with this library; if not, write to the
 * Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
 * Boston, MA  02110-1301, USA.
 *
 * Written by Luis Campo Giralte <luis.camp0.2009@gmail.com>
 *
 */
#include "DNSProtocol.h"
#include <iomanip> // setw

namespace aiengine {

DNSProtocol::DNSProtocol():
	Protocol("DNS", IPPROTO_UDP) {}

DNSProtocol::~DNSProtocol() {

	anomaly_.reset();
}

bool DNSProtocol::check(const Packet &packet) {

	// I dont like this idea of ports but...
	if (((packet.getSourcePort() == 53)||(packet.getDestinationPort() == 53))
		or((packet.getSourcePort() == 5353)||(packet.getDestinationPort() == 5353))) {
		if (packet.getLength() >= header_size) {
			setHeader(packet.getPayload());

			++total_valid_packets_;
			return true;
		}
	}
	++total_invalid_packets_;
	return false;
}

void DNSProtocol::setDynamicAllocatedMemory(bool value) {

	info_cache_->setDynamicAllocatedMemory(value);
	domain_cache_->setDynamicAllocatedMemory(value);
}

bool DNSProtocol::isDynamicAllocatedMemory() const {

	return info_cache_->isDynamicAllocatedMemory();
}

uint64_t DNSProtocol::getCurrentUseMemory() const {

        uint64_t mem = sizeof(DNSProtocol);

        mem += info_cache_->getCurrentUseMemory();
        mem += domain_cache_->getCurrentUseMemory();

        return mem;
}

uint64_t DNSProtocol::getAllocatedMemory() const {

        uint64_t mem = sizeof(DNSProtocol);

        mem += domain_cache_->getAllocatedMemory();
        mem += info_cache_->getAllocatedMemory();

        return mem;
}

uint64_t DNSProtocol::getTotalAllocatedMemory() const {

        uint64_t mem = getAllocatedMemory();

        mem += compute_memory_used_by_maps();

	return mem;
}

uint64_t DNSProtocol::compute_memory_used_by_maps() const {

	uint64_t bytes = 0;
	// Compute the size of the strings used as keys on the map
	std::for_each (domain_map_.begin(), domain_map_.end(), [&bytes] (PairStringCacheHits const &dt) {
		bytes += dt.first.size();
	});
	return bytes;
}

uint32_t DNSProtocol::getTotalCacheMisses() const {

	uint32_t miss = 0;

	miss = info_cache_->getTotalFails();
	miss += domain_cache_->getTotalFails();

	return miss;
}

void DNSProtocol::releaseCache() {

	if (FlowManagerPtr fm = flow_mng_.lock(); fm) {
		auto ft = fm->getFlowTable();

		std::ostringstream msg;
		msg << "Releasing " << name() << " cache";

		infoMessage(msg.str());

		uint64_t total_cache_bytes_released = compute_memory_used_by_maps();
		uint64_t total_bytes_released_by_flows = 0;
		uint64_t total_cache_save_bytes = 0;
		uint32_t release_flows = 0;
		uint32_t release_doms = domain_map_.size();

		for (auto &flow: ft) {
			if (SharedPointer<DNSInfo> info = flow->getDNSInfo(); info) {
				total_bytes_released_by_flows += sizeof(info);

				++release_flows;
				flow->layer7info.reset();
				info_cache_->release(info);
			}
		}
                // Some entries can be still on the maps and needs to be
                // retrieve to their existing caches
                for (auto &entry: domain_map_) {
			total_cache_save_bytes += entry.second.sc->size() * (entry.second.hits - 1);
                        releaseStringToCache(domain_cache_, entry.second.sc);
		}
		domain_map_.clear();

		data_time_ = boost::posix_time::microsec_clock::local_time();

		msg.str("");
		msg << "Release " << release_doms << " domains, " << release_flows << " flows";
		computeMemoryUtilization(msg, total_cache_bytes_released, total_bytes_released_by_flows, total_cache_save_bytes);
		infoMessage(msg.str());
	}
}

void DNSProtocol::releaseFlowInfo(Flow *flow) {

	if (auto info = flow->getDNSInfo(); info)
		info_cache_->release(info);
}

void DNSProtocol::attach_dns_to_flow(DNSInfo *info, boost::string_ref &domain, uint16_t qtype) {

	// Update the query type independently if there is memory for the name
	info->setQueryType(qtype);

	// Check if the domain is cache to attach
	if (StringMap::iterator it = domain_map_.find(domain); it != domain_map_.end()) {
		++(it->second).hits;
		info->name = (it->second).sc;
	} else {
		if (SharedPointer<StringCache> name = domain_cache_->acquire(); name) {
			name->name(domain.data(), domain.length());
                        info->name = name;
                        domain_map_.insert(std::make_pair(name->name(), name));
                }
	}
}

void DNSProtocol::processFlow(Flow *flow) {

	CPUCycle cycles(&total_cpu_cycles_);
	int length = flow->packet->getLength();
	total_bytes_ += length;
	++total_packets_;

	current_flow_ = flow;
	current_length_ = length;
	current_offset_ = 0;

	if (length >= header_size) { // Minimum header size consider
		setHeader(flow->packet->getPayload());

        	SharedPointer<DNSInfo> info = flow->getDNSInfo();
        	if (!info) {
                	if (info = info_cache_->acquire(); !info) {
				logFailCache(info_cache_->name(), flow);
				return;
                	}
                	flow->layer7info = info;
        	}

		if (info->isBanned() == false) {
			current_offset_ += header_size;

			if (header_->qr == 0) { // Query
				++total_queries_;
				if (ntohs(header_->questions) > 0)
					handle_standard_query(info.get(), length - header_size);
			} else { // Responses
				++total_responses_;
				handle_standard_response(info.get(), length - header_size);
			}
		}
	} else {
		++total_events_;
               	if (flow->getPacketAnomaly() == PacketAnomalyType::NONE)
               		flow->setPacketAnomaly(PacketAnomalyType::DNS_BOGUS_HEADER);

                anomaly_->incAnomaly(PacketAnomalyType::DNS_BOGUS_HEADER);
	}
	return;
}

int DNSProtocol::parse_query_name(Flow *flow, int length) {

	int offset = extract_domain_name(&header_->data[0], length);

	if (offset >= MAX_DNS_BUFFER_NAME) {
		++total_events_;
        	if (flow->getPacketAnomaly() == PacketAnomalyType::NONE)
                	flow->setPacketAnomaly(PacketAnomalyType::DNS_LONG_NAME);

                anomaly_->incAnomaly(PacketAnomalyType::DNS_LONG_NAME);
	}
	return offset;
}

// INFO: http://www.tcpipguide.com/free/t_DNSNameNotationandMessageCompressionTechnique.htm
int DNSProtocol::extract_domain_name(const uint8_t *ptr, int length) {
        int offset = 1;
        int8_t next = (int8_t)ptr[0];
	int max_length = MAX_DNS_BUFFER_NAME - 1;

	if (length < max_length)
		max_length = length;

	dns_buffer_name_[0] = '0';

	while ((offset < max_length) and (next > 0)) {
		if (next + offset < max_length) {
			std::memcpy(&dns_buffer_name_[offset - 1], &ptr[offset], next);
			offset += next + 1;
			next = (int8_t)ptr[offset - 1];
			if (next > 0 ) {
				dns_buffer_name_[offset - 2] = '.';
			}
		} else { // There is buffer for copy but the name is too long
			int left = max_length - offset;

			std::memcpy(&dns_buffer_name_[offset - 1], &ptr[offset], left + 1);
			offset = max_length + 2;
			break;
		}
	}

	if (offset > 1)
	    	-- offset;

	return offset;
}

bool DNSProtocol::parse_response_answer(DNSInfo *info, const uint8_t *ptr, int length, int answers) {

	int total_block = length - sizeof(dns_common_resource_record);;
	int coff = 0;

	info->ttl(std::numeric_limits<uint32_t>::max());

        // Extract the IP addresses or CNAME and store on the DNSDomain just when the domain have been matched
	for (int i = 0; i < answers && coff < total_block; ++i) {
		int off = 0;
		if ((ptr[0] & 0xC0) == 0) { // is not a pointer
			off = (int)ptr[0] + 1;
		}
		const dns_common_resource_record *ans = reinterpret_cast <const dns_common_resource_record*> (&ptr[off]);
		uint16_t block_length = ntohs(ans->length);
                uint16_t type = ntohs(ans->type);
                uint16_t class_type = ntohs(ans->class_type);
		uint32_t ttl = ntohl(ans->ttl);

		coff += block_length;
		current_offset_ += sizeof(dns_common_resource_record);

		if (block_length > (current_length_ - current_offset_)) {
			// The block have more length than the packet
			return true;
		}

		// We just store the lowest ttl of the responses
		if (ttl < info->ttl())
			info->ttl(ttl);

		if (class_type == 0x0001) { // class IN
                        if ((type == 0x0001)and(block_length == 4)) { // IPv4
                                uint32_t ipv4addr =  ((ans->data[3] << 24) + (ans->data[2] << 16) + (ans->data[1] << 8) + ans->data[0]);
                                in_addr a;

                                a.s_addr = ipv4addr;
                                info->addIPAddress(inet_ntoa(a));
                        } else if ((type == 0x001C)and(block_length == 16)) { // IPv6
                                char ipv6addr[INET6_ADDRSTRLEN];
                                in6_addr *in6addr = (in6_addr*)&(ans->data[0]);

                                inet_ntop(AF_INET6, in6addr, ipv6addr, INET6_ADDRSTRLEN);

                                info->addIPAddress(ipv6addr);
                        } else if (type == 0x0005) { // CNAME
                                int value __attribute__((unused)) = extract_domain_name(&ans->data[0], block_length);

                                info->addName(&dns_buffer_name_[0]);
                        } else if (type == 0x0010) { // TXT record
				const dns_txt_record *txt = reinterpret_cast<const dns_txt_record*>(&ans->data[0]);
				char *data = (char*)&txt->data[0];

				if (txt->length < block_length) {

					// Increment the current offset
					current_offset_ += txt->length;

					if (block_length == txt->length + 1) { // there is only one txt
						boost::string_ref txt_record(data, (int)txt->length);

						info->addName(txt_record.data(), txt_record.length());
					} else { // there is a txt record split in more, just process one more
						std::string temp_data(data, (int)txt->length);// copy the first block

						// Points to the second txt data record
						txt = reinterpret_cast<const dns_txt_record*>(&ans->data[txt->length + 1]);
						data = (char*)&txt->data[0];

						// Verify that is not corrupted
						if (txt->length > (current_length_ - current_offset_)) {
							return true;
						}
						temp_data.append(data, (int)txt->length);
						info->addName(temp_data.c_str());
					}
				}
			}
		}
                // TODO: Check offset size lengths and possible anomalies
                ptr = &(ans->data[block_length]);
        }
	return false;
}

void DNSProtocol::handle_standard_query(DNSInfo *info, int length) {
	boost::string_ref domain;
	int offset = parse_query_name(current_flow_, length);

	boost::string_ref dns_name(dns_buffer_name_, offset);

	if (offset == 1) { // There is no name, a root record
		offset = 0;
		domain = "<Root>";
	} else {
		domain = dns_name.substr(0, offset - 1);
	}

#ifdef DEBUG
	std::cout << __FILE__ << ":" << __func__ << ":length:" << length << " name:" << domain << " offset:" << offset << std::endl;
#endif
	// Check if the payload is malformed
	// The offset + 4 is because at the end of the domain 4 bytes should be present
	if (offset + 4 > length) {
		++total_events_;
               	if (current_flow_->getPacketAnomaly() == PacketAnomalyType::NONE)
               		current_flow_->setPacketAnomaly(PacketAnomalyType::DNS_BOGUS_HEADER);

               	anomaly_->incAnomaly(current_flow_, PacketAnomalyType::DNS_BOGUS_HEADER);
		return;
	}

	uint16_t qtype = ntohs((header_->data[offset + 2] << 8) + header_->data[offset + 1]);

	update_query_types(qtype);

	info->txid(getTransactionId());

	if (domain.length() > 0) { // The domain is valid
		if (ban_domain_mng_) {
			if (auto domain_candidate = ban_domain_mng_->getDomainName(domain); domain_candidate) {
				info->setIsBanned(true);
				++total_ban_queries_;
				return;
			}
		}

		++total_allow_queries_;

		attach_dns_to_flow(info, domain, qtype);
	}
}

void DNSProtocol::handle_standard_response(DNSInfo *info, int length) {
	boost::string_ref domain;

       	int offset = parse_query_name(current_flow_, length);

#ifdef DEBUG
	std::cout << __FILE__ << ":" << __func__ << ":no name attached, length:" << length << " offset:" << offset;
	std::cout << " questions:" << (int)ntohs(header_->questions) << " ans:" << ntohs(header_->answers) << std::endl;
#endif
	// Update the response of the dns and update the counters
	info->response_code(header_->rcode);
	update_response_code(header_->rcode);

	if (ntohs((header_->answers) > 0)or(ntohs(header_->questions) > 0)) {
		if (SharedPointer<StringCache> name = info->name; !name) {
			// Check if the payload is malformed
			// The offset + 4 is because at the end of the domain 4 bytes should be present
			if (offset + 4 > length) {
				++total_events_;
				if (current_flow_->getPacketAnomaly() == PacketAnomalyType::NONE)
					current_flow_->setPacketAnomaly(PacketAnomalyType::DNS_BOGUS_HEADER);

				anomaly_->incAnomaly(current_flow_, PacketAnomalyType::DNS_BOGUS_HEADER);
			} else {
				// There is no name attached so lets try to extract from the response

				boost::string_ref dns_name(dns_buffer_name_, offset);

				if (offset == 1) { // There is no name, a root record
					offset = 0;
					domain = "<Root>";
				} else {
					domain = dns_name.substr(0, offset - 1);
				}

				uint16_t qtype = ntohs((header_->data[offset + 2] << 8) + header_->data[offset + 1]);

				update_query_types(qtype);

				attach_dns_to_flow(info, domain, qtype);
			}
		} else {
			domain = name->name();
		}
#ifdef DEBUG
		std::cout << __FILE__ << ":" << __func__ << ":domain:" << domain << " length:" << length << " offset:" << offset << std::endl;
#endif

		// Check if the DNSProtocol have a DomainNameManager attached for match domains
		if (domain_mng_) {
			if (auto domain_candidate = domain_mng_->getDomainName(domain); domain_candidate) {
				++total_events_;
				// Need to increase by 4 the generate offset due to the type and class dns fields
				offset = offset + 5;
				uint16_t answers = ntohs(header_->answers);
				const uint8_t *ptr = &(header_->data[offset]);
				current_offset_ += offset;

				if (bool bogus = parse_response_answer(info, ptr, length - offset,  answers); bogus) {
					++total_events_;
					if (current_flow_->getPacketAnomaly() == PacketAnomalyType::NONE)
						current_flow_->setPacketAnomaly(PacketAnomalyType::DNS_BOGUS_HEADER);

					anomaly_->incAnomaly(current_flow_, PacketAnomalyType::DNS_BOGUS_HEADER);
				}
				info->matched_domain_name = domain_candidate;
#if defined(BINDING)
				current_flow_->packet->setForceAdaptorWrite(true); // The udp layer will call the databaseAdaptor update method
				if (domain_candidate->call.haveCallback())
					domain_candidate->call.executeCallback(current_flow_);
#endif
			}
		}
	}
}

void DNSProtocol::update_query_types(uint16_t type) {

	if (type == static_cast<uint16_t>(DNSQueryTypes::DNS_TYPE_A))
		++total_dns_type_a_;
	else if (type == static_cast<uint16_t>(DNSQueryTypes::DNS_TYPE_NS))
		++total_dns_type_ns_;
	else if (type == static_cast<uint16_t>(DNSQueryTypes::DNS_TYPE_CNAME))
		++total_dns_type_cname_;
	else if (type == static_cast<uint16_t>(DNSQueryTypes::DNS_TYPE_SOA))
		++total_dns_type_soa_;
	else if (type == static_cast<uint16_t>(DNSQueryTypes::DNS_TYPE_PTR))
		++total_dns_type_ptr_;
	else if (type == static_cast<uint16_t>(DNSQueryTypes::DNS_TYPE_MX))
		++total_dns_type_mx_;
	else if (type == static_cast<uint16_t>(DNSQueryTypes::DNS_TYPE_TXT))
		++total_dns_type_txt_;
	else if (type == static_cast<uint16_t>(DNSQueryTypes::DNS_TYPE_AAAA))
		++total_dns_type_aaaa_;
	else if (type == static_cast<uint16_t>(DNSQueryTypes::DNS_TYPE_LOC))
		++total_dns_type_loc_;
	else if (type == static_cast<uint16_t>(DNSQueryTypes::DNS_TYPE_SRV))
		++total_dns_type_srv_;
	else if (type == static_cast<uint16_t>(DNSQueryTypes::DNS_TYPE_DS))
		++total_dns_type_ds_;
	else if (type == static_cast<uint16_t>(DNSQueryTypes::DNS_TYPE_SSHFP))
		++total_dns_type_sshfp_;
	else if (type == static_cast<uint16_t>(DNSQueryTypes::DNS_TYPE_DNSKEY))
		++total_dns_type_dnskey_;
	else if (type == static_cast<uint16_t>(DNSQueryTypes::DNS_TYPE_TLSA))
		++total_dns_type_tlsa_;
	else if (type == static_cast<uint16_t>(DNSQueryTypes::DNS_TYPE_HTTPS))
		++total_dns_type_https_;
	else if (type == static_cast<uint16_t>(DNSQueryTypes::DNS_TYPE_IXFR))
		++total_dns_type_ixfr_;
	else if (type == static_cast<uint16_t>(DNSQueryTypes::DNS_TYPE_ANY))
		++total_dns_type_any_;
	else {
		++total_dns_type_others_;
		// std::cout << *current_flow_ << std::endl;
	}

}

void DNSProtocol::update_response_code(uint16_t rcode) {

	if (rcode == static_cast<uint16_t>(DNSResponseCodes::DNS_NOERROR))
		++total_dns_response_noerror_;
	else if (rcode == static_cast<uint16_t>(DNSResponseCodes::DNS_FORMERR))
		++total_dns_response_formerr_;
	else if (rcode == static_cast<uint16_t>(DNSResponseCodes::DNS_SERVFAIL))
		++total_dns_response_servfail_;
	else if (rcode == static_cast<uint16_t>(DNSResponseCodes::DNS_NXDOMAIN))
		++total_dns_response_nxdomain_;
	else if (rcode == static_cast<uint16_t>(DNSResponseCodes::DNS_NOTIMP))
		++total_dns_response_notimp_;
	else if (rcode == static_cast<uint16_t>(DNSResponseCodes::DNS_REFUSED))
		++total_dns_response_refused_;
	else if (rcode == static_cast<uint16_t>(DNSResponseCodes::DNS_YXDOMAIN))
		++total_dns_response_yxdomain_;
	else if (rcode == static_cast<uint16_t>(DNSResponseCodes::DNS_XRRSET))
		++total_dns_response_xrrset_;
	else if (rcode == static_cast<uint16_t>(DNSResponseCodes::DNS_NOTAUTH))
		++total_dns_response_notauth_;	
	else if (rcode == static_cast<uint16_t>(DNSResponseCodes::DNS_NOTZONE))
		++total_dns_response_notzone_;
	else
		++total_dns_response_others_;
}

void DNSProtocol::setDomainNameManager(const SharedPointer<DomainNameManager> &dm) {

	if (domain_mng_)
		domain_mng_->setPluggedToName("");

	if (dm) {
        	domain_mng_ = dm;
	        domain_mng_->setPluggedToName(name());
	} else {
		domain_mng_.reset();
	}
}

void DNSProtocol::statistics(std::basic_ostream<char> &out, int level, int32_t limit) const {

	showStatisticsHeader(out, level);

	if (level > 0) {
		if (ban_domain_mng_)
			out << "\t" << "Plugged banned domains from:" << ban_domain_mng_->name() << "\n";
		if (domain_mng_)
			out << "\t" << "Plugged domains from:" << domain_mng_->name() << "\n";
	}
	if (level > 3) {
		out << "\t" << "Total allow queries:    " << std::setw(10) << total_allow_queries_ << "\n"
			<< "\t" << "Total banned queries:   " << std::setw(10) << total_ban_queries_ << "\n"
			<< "\t" << "Total queries:          " << std::setw(10) << total_queries_ << "\n"
			<< "\t" << "Total responses:        " << std::setw(10) << total_responses_ << "\n"
			<< "\t" << "Total type A:           " << std::setw(10) << total_dns_type_a_ << "\n"
			<< "\t" << "Total type NS:          " << std::setw(10) << total_dns_type_ns_ << "\n"
			<< "\t" << "Total type CNAME:       " << std::setw(10) << total_dns_type_cname_ << "\n"
			<< "\t" << "Total type SOA:         " << std::setw(10) << total_dns_type_soa_ << "\n"
			<< "\t" << "Total type PTR:         " << std::setw(10) << total_dns_type_ptr_ << "\n"
			<< "\t" << "Total type MX:          " << std::setw(10) << total_dns_type_mx_ << "\n"
			<< "\t" << "Total type TXT:         " << std::setw(10) << total_dns_type_txt_ << "\n"
			<< "\t" << "Total type AAAA:        " << std::setw(10) << total_dns_type_aaaa_ << "\n"
			<< "\t" << "Total type LOC:         " << std::setw(10) << total_dns_type_loc_ << "\n"
			<< "\t" << "Total type SRV:         " << std::setw(10) << total_dns_type_srv_ << "\n"
			<< "\t" << "Total type DS:          " << std::setw(10) << total_dns_type_ds_ << "\n"
			<< "\t" << "Total type SSHFP:       " << std::setw(10) << total_dns_type_sshfp_ << "\n"
			<< "\t" << "Total type DNSKEY:      " << std::setw(10) << total_dns_type_dnskey_ << "\n"
			<< "\t" << "Total type TLSA:        " << std::setw(10) << total_dns_type_tlsa_ << "\n"
			<< "\t" << "Total type HTTPS:       " << std::setw(10) << total_dns_type_https_ << "\n"
			<< "\t" << "Total type IXFR:        " << std::setw(10) << total_dns_type_ixfr_ << "\n"
			<< "\t" << "Total type ANY:         " << std::setw(10) << total_dns_type_any_ << "\n"
			<< "\t" << "Total type others:      " << std::setw(10) << total_dns_type_others_ << "\n"
			<< "\t" << "Total rcode NOERROR:    " << std::setw(10) << total_dns_response_noerror_ << "\n"
			<< "\t" << "Total rcode FORMERR:    " << std::setw(10) << total_dns_response_formerr_ << "\n"
			<< "\t" << "Total rcode SERVFAIL:   " << std::setw(10) << total_dns_response_servfail_ << "\n"
			<< "\t" << "Total rcode NXDOMAIN:   " << std::setw(10) << total_dns_response_nxdomain_ << "\n"
			<< "\t" << "Total rcode NOTIMP:     " << std::setw(10) << total_dns_response_notimp_ << "\n"
			<< "\t" << "Total rcode REFUSED:    " << std::setw(10) << total_dns_response_refused_ << "\n"
			<< "\t" << "Total rcode YXDOMAIN:   " << std::setw(10) << total_dns_response_yxdomain_ << "\n"
			<< "\t" << "Total rcode XRRSET:     " << std::setw(10) << total_dns_response_xrrset_ << "\n"
			<< "\t" << "Total rcode NOTAUTH:    " << std::setw(10) << total_dns_response_notauth_ << "\n"
			<< "\t" << "Total rcode NOTZONE:    " << std::setw(10) << total_dns_response_notzone_ << "\n"
			<< "\t" << "Total rcode others:     " << std::setw(10) << total_dns_response_others_ << std::endl;
	}
	if ((level > 5)and(flow_forwarder_.lock()))
		flow_forwarder_.lock()->statistics(out);
	if (level > 3) {
		info_cache_->statistics(out);
		domain_cache_->statistics(out);
		if (level > 4)
			domain_map_.show(out, "\t", limit);
	}
}

void DNSProtocol::statistics(Json &out, int level) const {

	showStatisticsHeader(out, level);

        if (level > 3) {
		out["allow queries"] = total_allow_queries_;
		out["banned queries"] = total_ban_queries_;
		out["queries"] = total_queries_;
		out["responses"] = total_responses_;

		Json j;

		j["a"] = total_dns_type_a_;
		j["ns"] = total_dns_type_ns_;
		j["cname"] = total_dns_type_cname_;
		j["soa"] = total_dns_type_soa_;
		j["ptr"] = total_dns_type_ptr_;
		j["mx"] = total_dns_type_mx_;
		j["txt"] = total_dns_type_txt_;
		j["aaaa"] = total_dns_type_aaaa_;
		j["loc"] = total_dns_type_loc_;
		j["srv"] = total_dns_type_srv_;
		j["ds"] = total_dns_type_ds_;
		j["sshfp"] = total_dns_type_sshfp_;
		j["dnskey"] = total_dns_type_dnskey_;
		j["tlsa"] = total_dns_type_tlsa_;
		j["https"] = total_dns_type_https_;
		j["ixfr"] = total_dns_type_ixfr_;
		j["any"] = total_dns_type_any_;
		j["others"] = total_dns_type_others_;

		out["types"] = j;

		Json rj;

        	rj["noerror"] = total_dns_response_noerror_;
        	rj["formerr"] = total_dns_response_formerr_;
        	rj["servfail"] = total_dns_response_servfail_;
        	rj["nxdomain"] = total_dns_response_nxdomain_;
        	rj["notimp"] = total_dns_response_notimp_;
        	rj["refused"] = total_dns_response_refused_;
        	rj["yxdomain"] = total_dns_response_yxdomain_;
        	rj["xrrset"] = total_dns_response_xrrset_;
        	rj["notauth"] = total_dns_response_notauth_;
        	rj["notzone"] = total_dns_response_notzone_;
        	rj["others"] = total_dns_response_others_;

		out["rcodes"] = rj;
        }
}

void DNSProtocol::increaseAllocatedMemory(int value) {

	info_cache_->create(value);
	domain_cache_->create(value);
}

void DNSProtocol::decreaseAllocatedMemory(int value) {

	info_cache_->destroy(value);
	domain_cache_->destroy(value);
}

CounterMap DNSProtocol::getCounters() const {
	CounterMap cm;

        cm.addKeyValue("packets", total_packets_);
        cm.addKeyValue("bytes", total_bytes_);

        cm.addKeyValue("allow queries", total_allow_queries_);
        cm.addKeyValue("banned queries", total_ban_queries_);
        cm.addKeyValue("queries", total_queries_);
        cm.addKeyValue("responses", total_responses_);

        cm.addKeyValue("type A", total_dns_type_a_);
        cm.addKeyValue("type NS", total_dns_type_ns_);
        cm.addKeyValue("type CNAME", total_dns_type_cname_);
        cm.addKeyValue("type SOA", total_dns_type_soa_);
        cm.addKeyValue("type PTR", total_dns_type_ptr_);
        cm.addKeyValue("type MX", total_dns_type_mx_);
        cm.addKeyValue("type TXT", total_dns_type_txt_);
        cm.addKeyValue("type AAAA", total_dns_type_aaaa_);
        cm.addKeyValue("type LOC", total_dns_type_loc_);
        cm.addKeyValue("type SRV", total_dns_type_srv_);
        cm.addKeyValue("type DS", total_dns_type_ds_);
        cm.addKeyValue("type SSHFP", total_dns_type_sshfp_);
        cm.addKeyValue("type DNSKEY", total_dns_type_dnskey_);
        cm.addKeyValue("type TLSA", total_dns_type_tlsa_);
        cm.addKeyValue("type HTTPS", total_dns_type_https_);
        cm.addKeyValue("type IXFR", total_dns_type_ixfr_);
        cm.addKeyValue("type ANY", total_dns_type_any_);
        cm.addKeyValue("type others", total_dns_type_others_);

	cm.addKeyValue("rcode noerror", total_dns_response_noerror_);
	cm.addKeyValue("rcode formerr", total_dns_response_formerr_);
	cm.addKeyValue("rcode servfail", total_dns_response_servfail_);
	cm.addKeyValue("rcode nxdomain", total_dns_response_nxdomain_);
	cm.addKeyValue("rcode notimp", total_dns_response_notimp_);
	cm.addKeyValue("rcode refused", total_dns_response_refused_);
	cm.addKeyValue("rcode yxdomain", total_dns_response_yxdomain_);
	cm.addKeyValue("rcode xrrset", total_dns_response_xrrset_);
	cm.addKeyValue("rcode notauth", total_dns_response_notauth_);
	cm.addKeyValue("rcode notzone", total_dns_response_notzone_);
	cm.addKeyValue("rcode others", total_dns_response_others_);

        return cm;
}

#if defined(PYTHON_BINDING) || defined(RUBY_BINDING)
#if defined(PYTHON_BINDING)
boost::python::dict DNSProtocol::getCacheData(const std::string &name) const {
#elif defined(RUBY_BINDING)
VALUE DNSProtocol::getCacheData(const std::string &name) const {
#endif
	return addMapToHash(domain_map_);
}

#if defined(PYTHON_BINDING)
SharedPointer<Cache<StringCache>> DNSProtocol::getCache(const std::string &name) {

        if (boost::iequals(name, "domain"))
                return domain_cache_;

        return nullptr;
}

#endif

#endif

void DNSProtocol::statistics(Json &out, const std::string &map_name, int32_t limit) const {

        if (boost::iequals(map_name, "domains")) {
		domain_map_.show(out, limit);
        }
}

void DNSProtocol::resetCounters() {

	reset();

        total_allow_queries_ = 0;
        total_ban_queries_ = 0;
        total_queries_ = 0;
        total_responses_ = 0;
        total_events_ = 0;
        total_dns_type_a_ = 0;
        total_dns_type_ns_ = 0;
        total_dns_type_cname_ = 0;
        total_dns_type_soa_ = 0;
        total_dns_type_ptr_ = 0;
        total_dns_type_mx_ = 0;
        total_dns_type_txt_ = 0;
        total_dns_type_aaaa_ = 0;
        total_dns_type_loc_ = 0;
        total_dns_type_srv_ = 0;
        total_dns_type_ds_ = 0;
        total_dns_type_sshfp_ = 0;
        total_dns_type_dnskey_ = 0;
        total_dns_type_https_ = 0;
        total_dns_type_ixfr_ = 0;
        total_dns_type_any_ = 0;
        total_dns_type_others_ = 0;
        total_dns_response_noerror_ = 0;
        total_dns_response_formerr_ = 0;
        total_dns_response_servfail_ = 0;
        total_dns_response_nxdomain_ = 0;
        total_dns_response_notimp_ = 0;
        total_dns_response_refused_ = 0;
        total_dns_response_yxdomain_ = 0;
        total_dns_response_xrrset_ = 0;
        total_dns_response_notauth_ = 0;
        total_dns_response_notzone_ = 0;
        total_dns_response_others_ = 0;
}

} // namespace aiengine

